﻿//
// Created by xjchen on 12/20/22.
//

#ifndef MIDDLEWARE_CUDAMODULE_H
#define MIDDLEWARE_CUDAMODULE_H


#include "TaskModule.h"

#if defined(SUPPORT_CUDA_GPUS)

#include <cuda_runtime.h>
class CudaCtxMgr;
class DeviceCtx;

#endif

class CudaModule : public TaskModule {
public:
    CudaModule();

    ~CudaModule();

    /* 不再需要实现onCompute() */
    void onCompute(buffer_table_t *input, buffer_table_t *output) override;

#if defined(SUPPORT_CUDA_GPUS)

    /* 指定该CudaModule对应的GPU设备，暂不支持一个模块的算法运行在多个GPU设备上。
     * 由于设备切换比较耗时，该方法的目的是为了利用设备状态的线程保护功能，从而避免频繁调用cudaSetDevice()。
     * deviceID默认值为0，如果是其它值，建议在构造函数中变更。
     * */
    void setGPUDevice(int deviceID);

    /* 向指定stream推入一系列CUDA操作来实现GPU算法(如启动kernel，cudaMemcpyAsync，或者调用算法库的API等)。
     * 这些操作必须绑定到指定的stream上。
     * input和output对象驻留在host内存，但数据buffer(即buffer_node_t的data字段所指向的buffer)则被分配在对应GPU设备的内存上。
     * 由于不同CudaModule可并发运行，仍使用双buffer策略来降低不同模块间的同步等待时间。
     * 如果某端口连接的是同一GPU设备的CudaMoudle，则该连线不存在额外的数据拷贝，即直接引用数据指针。
     * 如果某端口连接另一GPU设备的CudaModule，或TaskModule，或DeviceReader，或DeviceWriter，则针对该连线中间件会做额外的数据拷贝。
     * 如果GPU不同设备间支持P2P拷贝，则优先使用P2P拷贝。
     * 支持积累，支持屏蔽输出端口，支持设置下次启动模式。
     * 可调用所有标记为USER_API的方法，中间件会在计算完成后做相应的状态处理。
     * 该方法仅在host端被调用，需要立即返回，不能被阻塞，不能做任何耗时操作。
     * */
    virtual void onIssueStreamedCudaOperations(buffer_table_t *input, buffer_table_t *output, cudaStream_t stream) = 0;

    int getGPUDevice();

private:
    friend class CudaCtxMgr;
    friend class DeviceCtx;
    void prev();
    void post();
    int _deviceID = 0;
    cudaStream_t _stream = 0;

#endif
};


#endif //MIDDLEWARE_CUDAMODULE_H
